import socketserver
import struct
import pickle
import threading
import os
import json
import asyncio
import websockets
USERS_FILE = "users.json"
USAGE_FILE = 'user_usage_data.json'

def load_users():
    if os.path.exists(USERS_FILE):
        with open(USERS_FILE, "r") as f:
            return json.load(f)
    return {}
def save_users(users):
    with open(USERS_FILE, "w") as f:
        json.dump(users, f, indent=2)
class DataStorage:

    def __init__(self, filepath: str):
        self.filepath = filepath
        self.today = '4.30'
        self.user_data = {}
        self._load_data()


    def _load_data(self):
        with open(self.filepath, 'r') as f:
            raw = json.load(f)

        for user, entries in raw.items():
            if not isinstance(entries, list) or len(entries) < 2:
                continue
            user_id = entries[0] if isinstance(entries[0], str) else None
            daily_entries = entries[1:]
            parsed_days = []

            for item in daily_entries:
                if not (isinstance(item, list) and len(item) == 2):
                    continue
                date, usage = item

                if date == self.today:
                    # detailed record for today
                    record = {
                        'date':      date,
                        'TikTok':    round(usage.get('TikTok', 0), 1),
                        'Youtube':   round(usage.get('Youtube', 0), 1),
                        'Instagram': round(usage.get('Instagram', 0), 1),
                        'Total':     round(usage.get('Total', 0), 1),
                    }
                else:
                    # simplified record for past days
                    video_time = (
                        usage.get('TikTok', 0) +
                        usage.get('Youtube', 0) +
                        usage.get('Instagram', 0)
                    )
                    record = {
                        'date':      date,
                        'videotime': round(video_time, 1),
                        'totaltime': round(usage.get('Total', 0), 1),
                    }

                parsed_days.append((date, record))

            self.user_data[user] = {
                'id':    user_id,
                'daily': parsed_days,

            }


    def get_user_data(self):
        return self.user_data

class RequestHandler:
    users = {}
    posts = []
    tokens = {}
    connected_clients = set()
    data_storage = DataStorage("user_usage_data.json")
    UsersLock = threading.Lock()
    #  TCP接口
    TcpCall = dict(
        LOGIN=(lambda self, *args: self.login(*args)),
        GET_USER_DATA=(lambda self, *args: self.data_storage.get_user_data()),
    )
    # WebSocket接口
    WsCall = dict(
        GET_POSTS=lambda cls: cls.posts,
        GET_TOKEN=lambda cls, username: cls.tokens.get(username, 0),
        CREATE_POST=lambda cls, data: cls.create_post(data),
        SET_TOKEN=lambda cls, data: cls.set_token(data)
    )
    def handle(self):
        SizeStruct = struct.Struct("!I")
        size_data = b''
        while len(size_data) < SizeStruct.size:
            more = self.rfile.read(SizeStruct.size - len(size_data))
            if not more:
                return  # 连接关闭
            size_data += more
        size = SizeStruct.unpack(size_data)[0]

        data_bytes = b''
        while len(data_bytes) < size:
            more = self.rfile.read(size - len(data_bytes))
            if not more:
                return  # 连接关闭
            data_bytes += more

        data = pickle.loads(data_bytes)
        command = data[0]
        args = data[1:]

        try:
            with RequestHandler.UsersLock:
                function = RequestHandler.TcpCall.get(command)
            if function is None:
                reply = (False, "Invalid command")
            else:
                reply = function(self, *args)
        except Exception as e:
            reply = (False, f"Server error: {str(e)}")

        reply_data = pickle.dumps(reply, protocol=3)
        self.wfile.write(SizeStruct.pack(len(reply_data)))
        self.wfile.write(reply_data)

    def login(self,username,password):
        # SizeStruct = struct.Struct("!I")
        with RequestHandler.UsersLock:
            if username in self.users:
                return False, "Username already exists"
            else:
                self.users[username] = password
                save_users(self.users)
                # data = {
                #     'username': username,
                #     'password': password,
                # }
                # reply_data = pickle.dumps(data, 3)
                # self.wfile.write(SizeStruct.pack(len(reply_data)))
                # self.wfile.write(reply_data
            return True, "User registered"
    def get_user_data(self):
        return self.user_data
    @classmethod
    async def handle_message(cls, websocket):
        cls.connected_clients.add(websocket)
        try:
            async for message in websocket:
                try:
                    data = pickle.loads(message)
                    command = data[0]
                    args = data[1:]

                    print(f"[SERVER] Received: {command} {args}")

                    if command in cls.WsCall:
                        handler = cls.WsCall[command]
                        if asyncio.iscoroutinefunction(handler):
                            result = await handler(cls, *args)
                        else:
                            result = handler(cls, *args)

                        await websocket.send(pickle.dumps((command, result), protocol=3))

                    else:
                        await websocket.send(pickle.dumps((command, (False, "Invalid command")), protocol=3))
                except Exception as e:
                    await websocket.send(pickle.dumps(("error", str(e)), protocol=3))
        finally:
            cls.connected_clients.discard(websocket)

    @classmethod
    def create_post(cls, data):
        if data['username'] in cls.users and data['password'] == cls.users[data['username']]:
            cls.posts.append(data)
            asyncio.create_task(cls.broadcast_all("GET_POSTS"))
            return True, "Post created"
        return False, "User does not exist"

    @classmethod
    def set_token(cls, data):
        if data['username'] in cls.users:
            cls.tokens[data['username']] = data['token']
            asyncio.create_task(cls.broadcast_all("GET_TOKEN", data['username']))
            return True, "Token set"
        return False, "User does not exist"

    @classmethod
    async def broadcast_all(cls, command, *args):
        reply = cls.Call[command](cls, *args)
        message = pickle.dumps((command, reply), protocol=3)
        await asyncio.gather(
            *(ws.send(message) for ws in cls.connected_clients if not ws.closed),
            return_exceptions=True
        )

class UserServer(socketserver.ThreadingMixIn, socketserver.TCPServer):
    pass
async def main():
    websocket_server = websockets.serve(RequestHandler.handle_message, "0.0.0.0", 8080)
    await websocket_server.__aenter__()
    print("WebSocket server started on ws://0.0.0.0:8080")

    # 启动 TCP 服务
    tcp_server = UserServer((" 192.168.0.198", 9999), RequestHandler)
    print("TCP server started on  192.168.0.198:9999")
    loop = asyncio.get_event_loop()
    tcp_future = loop.run_in_executor(None, tcp_server.serve_forever)

    # 等待两个服务同时运行
    await asyncio.gather(
        asyncio.Future(),  # 模拟永远运行
        tcp_future
    )

if __name__ == "__main__":
    asyncio.run(main())
